from typing import Tuple, List, Dict, Any
from rainforeLearn.gomoku.v2.environment.env import GomokuEnvironment
from rainforeLearn.gomoku.v2.agents.dqn_agent import GomokuDQNAgent
from rainforeLearn.gomoku.v2.train.run.game_state import GameState
from rainforeLearn.gomoku.v2.train.constants.train_constants import TrainConstants


class GameRunner:
    """游戏运行器 - 负责单个游戏的执行逻辑"""
    
    def __init__(self, env: GomokuEnvironment, agent: GomokuDQNAgent):
        self.env = env
        self.agent = agent
    
    def run_single_game(self, use_mcts: bool = False) -> Tuple[float, int, int, int]:
        """运行单个游戏，返回(奖励, 长度, 获胜者, 非法落子数)"""
        game_state = GameState()
        
        state = self.env.reset()
        
        while not self.env.game_over:
            action = self._select_action(state, use_mcts)
            
            if action == TrainConstants.INVALID_ACTION:
                break
            
            next_state, reward, done, info = self.env.step(action)
            
            # 更新游戏状态
            game_state.update(state, action, reward, next_state, done, info)
            state = next_state
        
        # 存储所有transitions
        self._store_transitions(game_state.transitions)
        
        return game_state.get_results(self.env.winner)
    
    def _select_action(self, state: Dict, use_mcts: bool) -> int:
        """选择动作"""
        if self.env.current_player == TrainConstants.AI_PLAYER:
            return self.agent.select_action(state, use_mcts=use_mcts)
        else:
            return self.agent.select_action(state, epsilon=0.0)
    
    def _store_transitions(self, transitions: List[Tuple]) -> None:
        """存储所有transitions到经验回放缓冲区"""
        for transition in transitions:
            self.agent.store_transition(*transition)
    
    def _state_to_array(self, state: Dict) -> Any:
        """将环境状态转换为数组"""
        return state['board'].flatten()